import requests

API_KEY = "7d766a92-7de7-4c14-1c90-1931628a54ac"


def gen_image(prompt,style):
    endpoint = "https://api.ttapi.io/midjourney/v1/imagine"

    headers = {
        "TT-API-KEY": API_KEY
    }
    data = {
        "prompt": f"{prompt} --ar 1:1",
        "model": "fast",
        "hookUrl": "",
        "getUImages": "true",
    }

    response = requests.post(endpoint, headers=headers, json=data)
    print(response.text)
    return response.json()["data"]["jobId"]


def get_image(jobId):
    print("正在获取图片地址")
    import requests

    endpoint = "https://api.ttapi.io/midjourney/v1/fetch"

    headers = {
        "TT-API-KEY": API_KEY
    }

    data = {
        "jobId": jobId
    }

    response = requests.post(endpoint, headers=headers, json=data)

    response_json = response.json()
    status = response_json["status"]
    if status == "SUCCESS":
        print(response_json)
        image_url = response_json["data"]["images"][0]
        return image_url

    if status == "ON_QUEUE":
        import time
        time.sleep(2)
        return get_image(jobId)


if __name__ == '__main__':
    job_id = gen_image("小猫","默认风格")
    print(job_id)
    url = get_image(job_id)
    print(url)
